package proxy

import (
	"errors"
	"fmt"
	"net"
	"time"
)

type ConnectionRelation struct {
	ProxyServer *ProxyServer
	SourceCon   net.Conn
	TargetCon   net.Conn
	UniqueId    string
}

func (ConnectionRelation *ConnectionRelation) SendToOtherConnection(conn net.Conn, data []byte) error {
	if conn.RemoteAddr() == ConnectionRelation.SourceCon.RemoteAddr() {
		return ConnectionRelation.SendToTargetCon(data)
	}
	if conn.RemoteAddr() == ConnectionRelation.TargetCon.RemoteAddr() {
		return ConnectionRelation.SendToSourceCon(data)
	}
	return errors.New("找不到连接")
}

func (ConnectionRelation *ConnectionRelation) GetToOtherConnection(conn net.Conn) (net.Conn, error) {
	if conn.RemoteAddr() == ConnectionRelation.SourceCon.RemoteAddr() {
		return ConnectionRelation.TargetCon, nil
	}
	if conn.RemoteAddr() == ConnectionRelation.TargetCon.RemoteAddr() {
		return ConnectionRelation.SourceCon, nil
	}
	return nil, errors.New("找不到连接")
}

func (ConnectionRelation *ConnectionRelation) SendToSourceCon(bytes []byte) error {
	retryTime := 1
	for ConnectionRelation.SourceCon == nil {
		retryTime++
		if retryTime > 5 {
			ConnectionRelation.TargetCon.Close()
			return errors.New("获取连接失败")
		}
		time.Sleep(time.Millisecond * 500)
	}
	_, err := ConnectionRelation.SourceCon.Write(bytes)
	if err != nil {
		ConnectionRelation.TargetCon.Close()
		fmt.Println(err)
		fmt.Println("发送失败")
		return err
	}
	return nil
}

func (ConnectionRelation *ConnectionRelation) SendToTargetCon(bytes []byte) error {
	retryTime := 1
	for ConnectionRelation.TargetCon == nil {
		retryTime++
		if retryTime > 5 {
			ConnectionRelation.SourceCon.Close()
			return errors.New("获取连接失败")
		}
		time.Sleep(time.Millisecond * 500)
	}
	_, err := ConnectionRelation.TargetCon.Write(bytes)
	if err != nil && ConnectionRelation.ProxyServer == nil {
		ConnectionRelation.SourceCon.Close()
		return err
	}
	for err != nil {
		fmt.Println("发送失败")
		conn, err := ConnectionRelation.ProxyServer.GetProxyConn(ConnectionRelation.UniqueId)
		if err != nil {
			fmt.Println("获取工作进程失败")
			ConnectionRelation.SourceCon.Close()
			break
		} else {
			ConnectionRelation.TargetCon = conn.conn
			_, err = ConnectionRelation.TargetCon.Write(bytes)
			if err == nil {
				break
			}
		}
	}
	return nil
}
